# coding:utf-8
"""
    对原始数据进行预处理，根据标注好的数据和原始数据进行转换为BIO类型的数据，保存为train.txt文件
        B：表示开头
        I：表示中间或结尾
        O：表示其他如标点符号和不重要字符
"""
import os
import json
from collections import Counter

# 切换工作目录, 获取源数据目录
os.chdir('..')


class DataProcess(object):
    def __init__(self):
        # 获取标签dict
        self.label_dict = json.load(open('./data/labels.json', 'r', encoding='utf-8'))
        self.tag2id = json.load(open('./data/tag2id.json', 'r', encoding='utf-8'))
        self.data_original_path = os.path.join(os.getcwd(), 'data_origin')
        self.train_file_path = os.path.join(os.getcwd(), 'data/train.txt')

    # 对每个标注好的文件进行读取，并根据原始数据进行打标签
    def transfer(self):
        # 获取所有文件名，其中源文件含有original，而标注好的文件没有
        all_file_names = os.walk(self.data_original_path)
        with open(self.train_file_path, 'w', encoding='utf-8') as fw:
            for root, dirs, file_names in all_file_names:
                for file_name in file_names:
                    if 'original' not in file_name:
                        continue
                    complete_file = file_name.replace('.txtoriginal', '')
                    com_file_path = root + '\\' + complete_file
                    re_dict = self.read_lable_text(com_file_path)
                    # print(re_dict)
                    or_file_path = root + '\\' + file_name
                    with open(or_file_path, 'r', encoding='utf-8') as f:
                        for line in f:
                            for index, s in enumerate(line.strip()):
                                char_label = re_dict.get(index, 'O')
                                fw.write(s + '\t' + char_label + '\n')
                        # print('写入文件成功')

    def read_lable_text(self, com_file_path):
        res_dict = {}
        with open(com_file_path, 'r', encoding='utf-8') as f:
            for line in f:
                # print(line)
                ll = line.strip().split('\t')
                start = int(ll[1])
                end = int(ll[2])
                label = ll[3]
                label_tag = self.label_dict.get(label)
                for i in range(start, end + 1):
                    if i == start:
                        tag = "B-" + label_tag
                    else:
                        tag = "I-" + label_tag
                    res_dict[i] = tag
        return res_dict


if __name__ == '__main__':
    data_process = DataProcess()
    data_process.transfer()
